//##########################################################################
//#                                                                        #
//#                       CLOUDCOMPARE PLUGIN: qEDL                        #
//#                                                                        #
//#  This program is free software; you can redistribute it and/or modify  #
//#  it under the terms of the GNU General Public License as published by  #
//#  the Free Software Foundation; version 2 of the License.               #
//#                                                                        #
//#  This program is distributed in the hope that it will be useful,       #
//#  but WITHOUT ANY WARRANTY; without even the implied warranty of        #
//#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         #
//#  GNU General Public License for more details.                          #
//#                                                                        #
//#          COPYRIGHT: EDF R&D / TELECOM ParisTech (ENST-TSI)             #
//#                                                                        #
//##########################################################################

#include "ccSSAOFilter.h"

#include <ccFrameBufferObject.h>
#include <ccBilateralFilter.h>
#include <ccShader.h>
#include <ccFBOUtils.h>

//RandomKit
#include <randomkit.h>

//system
#include <stdlib.h>
#include <math.h>
#include <vector>

#ifndef GL_RGBA32F
#define GL_RGBA32F 0x8814
#endif

#ifndef GL_RGB16F
#define GL_RGB16F 0x881B
#endif

ccSSAOFilter::ccSSAOFilter()
	: ccGlFilter("Screen Space Ambient Occlusion")
	, m_texReflect(0)
	, m_w(0)
	, m_h(0)
	, m_fbo(0)
	, m_shader(0)
{
	setParameters(/*N=*/32,/*Kz=*/500.0f,/*R=*/0.05f,/*F=*/50.0f);

	m_bilateralFilterEnabled = false;
	m_bilateralFilter        = 0;
	m_bilateralGHalfSize     = 2;
	m_bilateralGSigma        = 1.0f;
	m_bilateralGSigmaZ       = 0.4f;

	memset(m_ssao_neighbours, 0, sizeof(float) * 3*MAX_N);
	sampleSphere();
}

ccSSAOFilter::~ccSSAOFilter()
{
	reset();
}

ccGlFilter* ccSSAOFilter::clone() const
{
	ccSSAOFilter* filter = new ccSSAOFilter();

	//copy parameters
	filter->setParameters(m_N, m_Kz, m_R, m_F);

	return filter;
}

void ccSSAOFilter::reset()
{
	if (glIsTexture(m_texReflect))
		glDeleteTextures(1,&m_texReflect);
	m_texReflect = 0;

	if (m_fbo)
		delete m_fbo;
	m_fbo = 0;

	if (m_shader)
		delete m_shader;
	m_shader = 0;

	if (m_bilateralFilter)
		delete m_bilateralFilter;
	m_bilateralFilter = 0;
}

bool ccSSAOFilter::init(int width, int height, QString shadersPath, QString& error)
{
	return init(width,height,true,true,shadersPath,error);
}

bool ccSSAOFilter::init(int width,
						int height,
						bool enableBilateralFilter,
						bool useReflectTexture,
						QString shadersPath,
						QString& error,
						GLenum textureMinMagFilter /*= GL_LINEAR*/)
{
	//in case of reinit
	if (!m_fbo)
		m_fbo = new ccFrameBufferObject();
	if (!m_fbo->init(width,height))
	{
		error = "[SSAO] FrameBufferObject initialization failed!";
		reset();
		return false;
	}
	m_fbo->initTexture(0,GL_RGBA32F,GL_RGBA,GL_FLOAT,textureMinMagFilter);

	if (!m_shader)
	{
		m_shader = new ccShader();
		if (!m_shader->fromFile(shadersPath,"SSAO/ssao",error))
		{
			reset();
			return false;
		}
	}

	m_bilateralFilterEnabled = enableBilateralFilter;
	if (m_bilateralFilterEnabled)
	{
		if (!m_bilateralFilter)
			m_bilateralFilter = new ccBilateralFilter();
		if (!m_bilateralFilter->init(width,height,shadersPath,error))
		{
			delete m_bilateralFilter;
			m_bilateralFilter = 0;
			m_bilateralFilterEnabled = false;
		}
		else
		{
			m_bilateralFilter->useExistingViewport(true);
		}
	}
	else if (m_bilateralFilter)
	{
		delete m_bilateralFilter;
		m_bilateralFilter = 0;
	}

	m_w = width;
	m_h = height;

	if (useReflectTexture)
	{
		initReflectTexture();
	}
	else
	{
		if (m_texReflect != 0)
			glDeleteTextures(1,&m_texReflect);
		m_texReflect = 0;
	}

	return true;
}

void ccSSAOFilter::sampleSphere()
{
	// Initialize the sobol QRNG
	rk_sobol_state s;
	if (rk_sobol_init(3, &s, NULL, rk_sobol_Ldirections, NULL) != RK_SOBOL_OK)
	{
		return;
	}
	rk_sobol_randomshift(&s, NULL);

	//	draw in sphere
	float* ssao_neighbours = m_ssao_neighbours;

	for (unsigned n_in_sphere = 0; n_in_sphere < MAX_N; )
	{
		double x[5];
		rk_sobol_double(&s, x);
		
		double px = x[0]*2.0-1.0;
		double py = x[1]*2.0-1.0;
		double pz = x[2]*2.0-1.0;
		
		if ( px*px + py*py + pz*pz < 1.0 )
		{
			*ssao_neighbours++ = static_cast<float>(px);
			*ssao_neighbours++ = static_cast<float>(py);
			*ssao_neighbours++ = static_cast<float>(pz);

			++n_in_sphere;
		}
	}

	rk_sobol_free(&s);
}

void ccSSAOFilter::shade(GLuint texDepth, GLuint texColor, ViewportParameters& parameters)
{
	if (!m_fbo || !m_shader)
	{
		//ccLog::Warning("[ccSSAOFilter::shade] Internal error: structures not initialized!");
		return;
	}

	glPushAttrib(GL_ALL_ATTRIB_BITS);

	//we must use corner-based screen coordinates
	glMatrixMode(GL_PROJECTION);
	glPushMatrix();
	glLoadIdentity();
	glOrtho(0.0,(GLdouble)m_w,0.0,(GLdouble)m_h,0.0,1.0);
	glMatrixMode(GL_MODELVIEW);
	glPushMatrix();
	glLoadIdentity();

	m_fbo->start();
	m_fbo->setDrawBuffers1();

	m_shader->start();
	m_shader->setUniform1i("s2_Z",0);
	m_shader->setUniform1i("s2_R",1);
	m_shader->setUniform1i("s2_C",2);
	m_shader->setUniform1f("R",m_R);
	m_shader->setUniform1f("F",m_F);
	m_shader->setUniform1f("Kz",m_Kz);
	//m_shader->setUniform1i("N",N);
	m_shader->setUniform1i("B_REF",m_texReflect == 0 ? 0 : 1);
	m_shader->setTabUniform3fv("P",MAX_N,m_ssao_neighbours);

	glActiveTexture(GL_TEXTURE2);
	glEnable(GL_TEXTURE_2D);
	glBindTexture(GL_TEXTURE_2D,texColor);

	if (glIsTexture(m_texReflect))
	{
		glActiveTexture(GL_TEXTURE1);
		glEnable(GL_TEXTURE_2D);
		glBindTexture(GL_TEXTURE_2D,m_texReflect);
	}
	glActiveTexture(GL_TEXTURE0);

	ccFBOUtils::DisplayTexture2DCorner(texDepth,m_w,m_h);

	if (glIsTexture(m_texReflect))
	{
		glActiveTexture(GL_TEXTURE1);
		glBindTexture(GL_TEXTURE_2D,0);
		glDisable(GL_TEXTURE_2D);
	}

	glActiveTexture(GL_TEXTURE2);
	glBindTexture(GL_TEXTURE_2D,0);
	glDisable(GL_TEXTURE_2D);

	m_shader->stop();
	m_fbo->stop();

	if (m_bilateralFilter)
	{
		m_bilateralFilter->setParams(m_bilateralGHalfSize,m_bilateralGSigma,m_bilateralGSigmaZ);
		m_bilateralFilter->shade(texDepth,m_fbo->getColorTexture(0),parameters);
	}

	glMatrixMode(GL_PROJECTION);
	glPopMatrix();
	glMatrixMode(GL_MODELVIEW);
	glPopMatrix();

	glPopAttrib();
}

GLuint ccSSAOFilter::getTexture()
{
	if (m_bilateralFilter)
		return m_bilateralFilter->getTexture();

	return (m_fbo ? m_fbo->getColorTexture(0) : 0);
}

void ccSSAOFilter::setParameters(int _N, float _Kz, float _R, float _F)
{
	m_N  = _N;
	m_Kz = _Kz;
	m_R  = _R;
	m_F  = _F;
}

inline double frand()
{
	return static_cast<double>(rand()) / RAND_MAX;
}

void randomPointInSphere(double& vx, double& vy, double& vz)
{
	do
	{
		vx = frand();
		vy = frand();
		vz = frand();
	}
	while (vx*vx + vy*vy + vz*vz > 1.0);
}

void ccSSAOFilter::initReflectTexture()
{
	/***	INIT TEXTURE OF RELFECT VECTORS		***/
	/**		Fully random texture	**/
	int texSize = m_w*m_h;
	std::vector<float> reflectTex;
	try
	{
		reflectTex.resize(3*texSize);
	}
	catch (const std::bad_alloc&)
	{
		//not enough memory
		return;
	}

	for (int i=0; i<texSize; i++)
	{
		double x,y,z;
		randomPointInSphere(x,y,z);

		double norm = x*x+y*y+z*z;
		norm = (norm > 1.0e-12 ? 1.0/sqrt(norm) : 0.0);

		reflectTex[i*3  ] = static_cast<float>((1.0+x*norm)/2);
		reflectTex[i*3+1] = static_cast<float>((1.0+y*norm)/2);
		reflectTex[i*3+2] = static_cast<float>((1.0+z*norm)/2);
	}

	glEnable(GL_TEXTURE_2D);

	if (glIsTexture(m_texReflect))
		glDeleteTextures(1,&m_texReflect);

	glGenTextures(1,&m_texReflect);
	glBindTexture(GL_TEXTURE_2D,m_texReflect);

	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST );
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST );
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
	glTexImage2D(GL_TEXTURE_2D,0,GL_RGB16F,m_w,m_h,0,GL_RGB,GL_FLOAT,&reflectTex[0]);

	glDisable(GL_TEXTURE_2D);

	// According to Wikipedia, noise is made of 4*4 repeated tiles	to have only high frequency
}
